#include "ecs_app.hpp"

#include <iostream>
#include <stdexcept>
#include <array>
#include <vector>
#include <cassert>
#include <limits>
#include <numeric>
#include <algorithm>
#include <cmath>
#include <vulkan/vulkan.h>
#include <chrono>
#include <random>

#include "njm_utils.hpp"
#include "keyboard_movement_controller.hpp"
#include "njm_camera.hpp"
#include "njm_buffer.hpp"
#include "njm_descriptor.hpp"
#include "njm_game_obj.hpp"

#include "njm_ecs_frame_info.hpp"
#include "ecs_sys_importer.hpp"
#include "ecs_components.hpp"

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#include <glm/glm.hpp>
#include <glm/gtc/constants.hpp>

namespace njm
{
    // datas
    std::vector<glm::vec3> positions = {
        {-2.0f, -1.0f, 0.0f},
        {0.0f, -0.5f, 0.5f},
        {2.0f, 0.0f, 1.0f},
        {4.0f, 0.5f, 1.5f},
    };
    std::vector<glm::vec3> rotations = {
        {0.0f, 0.0f, 0.0f},
        {1.0f, 0.0f, 0.0f},
        {2.0f, 0.0f, 0.0f},
        {3.0f, 0.0f, 0.0f},
    };
    std::vector<glm::vec3> scales = {
        {1.0f, 1.0f, 1.0f},
        {1.0f, 1.0f, 1.0f},
        {1.0f, 1.0f, 1.0f},
        {1.0f, 1.0f, 1.0f},
    };
    std::vector<glm::vec3> forces = {
        {0.0f, -9.8f, 0.0f},
        {0.0f, -9.8f, 0.0f},
        {0.0f, -9.8f, 0.0f},
        {0.0f, -9.8f, 0.0f},
    };
    std::vector<glm::vec3> velocities = {
        {0.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 0.0f},
    };
    std::vector<glm::vec3> accelerations = {
        {0.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 0.0f},
    };
    std::vector<glm::vec4> colors = {
        {1.0f, 0.0f, 0.0f, 1.0f},
        {0.0f, 1.0f, 0.0f, 1.0f},
        {0.0f, 0.0f, 1.0f, 1.0f},
        {1.0f, 1.0f, 0.0f, 1.0f},
    };
    std::vector<std::string> tags = {
        "entity1",
        "entity2",
        "entity3",
        "entity4",
    };
    std::vector<NjmModel::id_t> modelIds = {
        1,
        2,
        3,
        1,
    };

    std::vector<Tag> getTags()
    {
        std::vector<Tag> rst;
        for (int i = 0; i < tags.size(); i++)
        {
            rst.push_back({tags[i]});
        }
        return rst;
    }
    std::vector<Model> getModels()
    {
        std::vector<Model> rst;
        for (int i = 0; i < modelIds.size(); i++)
        {
            rst.push_back({modelIds[i]});
        }
        return rst;
    }
    std::vector<Transform> getTransforms()
    {
        std::vector<Transform> transforms;
        for (int i = 0; i < positions.size(); i++)
        {
            transforms.push_back({positions[i], rotations[i], scales[i]});
        }
        return transforms;
    }
    std::vector<Gravity> getGravities()
    {
        std::vector<Gravity> gravities;
        for (int i = 0; i < forces.size(); i++)
        {
            gravities.push_back({forces[i]});
        }
        return gravities;
    }
    std::vector<RigidBody> getRigidBodies()
    {
        std::vector<RigidBody> rigidBodies;
        for (int i = 0; i < velocities.size(); i++)
        {
            rigidBodies.push_back({velocities[i], accelerations[i]});
        }
        return rigidBodies;
    }
    std::vector<Material> getMaterials()
    {
        std::vector<Material> materials;
        for (int i = 0; i < colors.size(); i++)
        {
            materials.push_back({colors[i]});
        }
        return materials;
    }

    // ECS APP
    ECSApp::ECSApp()
    {
        init();
    }
    ECSApp::~ECSApp()
    {
        mCoordinator->Clear();
    }
    void ECSApp::init()
    {
        // VULKAN
        globalPool = NjmDescriptorPool::Builder(njmDevice)
                         .setMaxSets(NjmSwapChain::MAX_FRAMES_IN_FLIGHT)
                         .addPoolSize(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, NjmSwapChain::MAX_FRAMES_IN_FLIGHT)
                         .addPoolSize(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, NjmSwapChain::MAX_FRAMES_IN_FLIGHT)
                         .build();
        loadImgs();
        loadModels();
        printAssetsInfo();
        //
        mCoordinator = std::make_unique<Coordinator>();

        mCoordinator->RegisterComponent<Tag>();
        mCoordinator->RegisterComponent<Model>();
        mCoordinator->RegisterComponent<Transform>();
        mCoordinator->RegisterComponent<Gravity>();
        mCoordinator->RegisterComponent<RigidBody>();
        mCoordinator->RegisterComponent<Material>();

        ComponentType tagType = mCoordinator->GetComponentType<Tag>();
        ComponentType modelType = mCoordinator->GetComponentType<Model>();
        ComponentType transformType = mCoordinator->GetComponentType<Transform>();
        ComponentType gravityType = mCoordinator->GetComponentType<Gravity>();
        ComponentType rigidBodyType = mCoordinator->GetComponentType<RigidBody>();
        ComponentType materialType = mCoordinator->GetComponentType<Material>();

        Signature tagSig = mCoordinator->GetComponentSignature<Tag>();
        Signature modelSig = mCoordinator->GetComponentSignature<Model>();
        Signature transformSig = mCoordinator->GetComponentSignature<Transform>();
        Signature gravitySig = mCoordinator->GetComponentSignature<Gravity>();
        Signature rigidBodySig = mCoordinator->GetComponentSignature<RigidBody>();
        Signature materialSig = mCoordinator->GetComponentSignature<Material>();
        Signature signature = tagSig | modelSig | transformSig | gravitySig | rigidBodySig | materialSig;

        //
        mCoordinator->RegisterSystem<PositionSystem>();
        mCoordinator->RegisterSystem<VelocitySystem>();
        mCoordinator->RegisterSystem<PhysicsSystem>();
        mCoordinator->RegisterSystem<EcsRenderSystem>();
        mCoordinator->RegisterSystem<TimeSystem>();

        mCoordinator->SetSystemSignature<PositionSystem>(transformSig);
        mCoordinator->SetSystemSignature<VelocitySystem>(rigidBodySig);
        mCoordinator->SetSystemSignature<PhysicsSystem>(signature);
        mCoordinator->SetSystemSignature<EcsRenderSystem>(modelSig | transformSig | materialSig);
        mCoordinator->SetSystemSignature<TimeSystem>(Signature());
        //
        std::default_random_engine generator;
        std::uniform_real_distribution<float> randPosition(-100.0f, 100.0f);
        std::uniform_real_distribution<float> randRotation(0.0f, 3.0f);
        std::uniform_real_distribution<float> randScale(3.0f, 5.0f);
        std::uniform_real_distribution<float> randGravity(-10.0f, -1.0f);
        std::uniform_real_distribution<float> randVelocity(-10.0f, 10.0f);
        std::uniform_real_distribution<float> randAcceleration(-10.0f, 10.0f);
        // float randData = randPosition(generator);

        std::vector<Tag> tags = getTags();
        std::vector<Model> models = getModels();
        std::vector<Transform> transforms = getTransforms();
        std::vector<Gravity> gravities = getGravities();
        std::vector<RigidBody> rigidBodies = getRigidBodies();
        std::vector<Material> materials = getMaterials();

        std::vector<Entity> entities = mCoordinator->CreateEntities(4);
        for (int i = 0; i < entities.size(); i++)
        {
            mCoordinator->AddComponent<Tag>(entities[i], tags[i]);
            mCoordinator->AddComponent<Model>(entities[i], models[i]);
            mCoordinator->AddComponent<Transform>(entities[i], transforms[i]);
            mCoordinator->AddComponent<Gravity>(entities[i], gravities[i]);
            mCoordinator->AddComponent<RigidBody>(entities[i], rigidBodies[i]);
            mCoordinator->AddComponent<Material>(entities[i], materials[i]);

            mCoordinator->AddEntityToSystem<PositionSystem>(entities[i]);
            mCoordinator->AddEntityToSystem<VelocitySystem>(entities[i]);
            mCoordinator->AddEntityToSystem<PhysicsSystem>(entities[i]);
            mCoordinator->AddEntityToSystem<EcsRenderSystem>(entities[i]);
            mCoordinator->AddEntityToSystem<TimeSystem>(entities[i]);
        }
    }

    void ECSApp::run()
    {
        std::cout << "------------ECSApp::run----------" << std::endl;
        //
        auto minOffsetAlignment = 1;
        std::vector<std::unique_ptr<NjmBuffer>> uboBuffers(NjmSwapChain::MAX_FRAMES_IN_FLIGHT);
        for (int i = 0; i < NjmSwapChain::MAX_FRAMES_IN_FLIGHT; i++)
        {
            uboBuffers[i] = std::make_unique<NjmBuffer>(
                njmDevice,
                sizeof(ECSGlobalUbo),
                1,
                VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
                minOffsetAlignment); // alignment also can be default value 1.
            uboBuffers[i]->map();
        }

        // create textures
        std::vector<std::unique_ptr<NjmTexture>> texs(NjmSwapChain::MAX_FRAMES_IN_FLIGHT);
        std::vector<TextureImgBinding> texImgBindings{
            {"tex1", "eg1.png"},
        };
        TextureImgBinding texImgBinding = texImgBindings[0];
        for (int i = 0; i < NjmSwapChain::MAX_FRAMES_IN_FLIGHT; i++)
        {
            auto &imgName = texImgBinding.imgName;
            auto &texName = texImgBinding.texName;
            assert(njmImgs.find(imgName) != njmImgs.end() && ("img not found : " + imgName).c_str());
            texs[i] = std::make_unique<NjmTexture>(njmDevice, texName, njmImgs.at(imgName));
        }

        // create descriptor set layouts and descriptor sets
        auto globalSetLayout = NjmDescriptorSetLayout::Builder(njmDevice)
                                   .addBinding(0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_ALL_GRAPHICS)
                                   .addBinding(1, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_FRAGMENT_BIT)
                                   .build();
        std::vector<VkDescriptorSet> globalDescriptorSets(NjmSwapChain::MAX_FRAMES_IN_FLIGHT);

        // update descriptor sets
        for (int i = 0; i < NjmSwapChain::MAX_FRAMES_IN_FLIGHT; i++)
        {
            auto bufferInfo = uboBuffers[i]->descriptorInfo();
            auto combinedImageInfo = texs[i]->getCombinedDescriptor();

            NjmDescriptorWriter(*globalSetLayout, *globalPool)
                .writeBuffer(0, &bufferInfo)
                .writeImage(1, &combinedImageInfo)
                .build(globalDescriptorSets[i]);
        }

        std::string SDR_BASE = "src/shaders/";
        std::string vertShaderPath = SDR_BASE + "ecs_shader_vert.spv";
        std::string fragShaderPath = SDR_BASE + "ecs_shader_frag.spv";
        auto trueRenderSys = std::make_shared<RenderSystem>(njmDevice,
                                                            njmRenderer.getSwapChainRenderPass(), globalSetLayout->getDescriptorSetLayout(),
                                                            vertShaderPath, fragShaderPath);

        KeyboardMovementController cameraController{};
        NjmGameObj viewerObject = NjmGameObj::createGameObj();
        viewerObject.transform.translation.z = -2.5f;
        NjmCamera camera{};
        // run loop
        bool quit = false;
        float timeFactor = 0.001f;
        auto currentTime = std::chrono::high_resolution_clock::now();
        while (!quit && !njmWindow.shouldClose())
        {
            std::cout << "-------------event looping------------ " << std::endl;
            glfwPollEvents();
            auto newTime = std::chrono::high_resolution_clock::now();
            float frameTime = std::chrono::duration<float, std::chrono::seconds::period>(newTime - currentTime).count();
            currentTime = newTime;
            frameTime = glm::min(frameTime, MAX_FRAME_TIME); // frametime is dt
            
            cameraController.moveInPlaneXZ(njmWindow.getGLFWwindow(), frameTime, viewerObject);
            camera.setViewYXZ(viewerObject.transform.translation, viewerObject.transform.rotation);
            float aspect = njmRenderer.getAspectRatio();
            camera.setPerspectiveProjection(glm::radians(50.f), aspect, 0.1f, 100.f);

            if (auto cmdBuf = njmRenderer.beginFrame())
            {
                int frameIndex = njmRenderer.getFrameIndex();
                ECSFrameInfo frameInfo{
                    frameIndex, frameTime,
                    cmdBuf, camera,
                    globalDescriptorSets[frameIndex], mCoordinator, njmModels};
                // update ubo data
                ECSGlobalUbo ubo{};
                ubo.view = camera.getView();
                ubo.projection = camera.getProjection();
                ubo.inverseView = camera.getInverseView();
                uboBuffers[frameIndex]->writeToBuffer(&ubo);
                uboBuffers[frameIndex]->flush();
                // render
                njmRenderer.beginSwapChainRenderPass(cmdBuf);
                trueRenderSys->render(frameInfo);
                njmRenderer.endSwapChainRenderPass(cmdBuf);
                njmRenderer.endFrame();
            }
            mCoordinator->Update(frameTime * timeFactor);
        }
        vkDeviceWaitIdle(njmDevice.device());
    }

    //
    std::unique_ptr<NjmModel> ecsCreateCubeModel(NjmDevice &device, glm::vec3 offset)
    {
        NjmModel::Builder builder{};
        // builder.loadModelFromFile(device, "models/cube.obj");
        // auto cubeModel = std::make_unique<NjmModel>(device, builder);
        // return cubeModel;

        // 立方体顶点数据（8个顶点）
        std::vector<NjmModel::Vertex> cubeVertices = {
            // Front face
            {{-0.5f, -0.5f, 0.5f}, {1.0f, 0.0f, 0.0f}}, // 红色
            {{0.5f, -0.5f, 0.5f}, {0.0f, 1.0f, 0.0f}},  // 绿色
            {{-0.5f, 0.5f, 0.5f}, {0.0f, 0.0f, 1.0f}},  // 蓝色
            {{0.5f, 0.5f, 0.5f}, {1.0f, 1.0f, 0.0f}},   // 黄色
            // Back face
            {{-0.5f, -0.5f, -0.5f}, {1.0f, 0.0f, 1.0f}}, // 紫色
            {{0.5f, -0.5f, -0.5f}, {0.0f, 1.0f, 1.0f}},  // 青色
            {{-0.5f, 0.5f, -0.5f}, {1.0f, 1.0f, 1.0f}},  // 白色
            {{0.5f, 0.5f, -0.5f}, {0.0f, 0.0f, 0.0f}}    // 黑色
        };
        // 立方体索引数据（36个索引，6个面×2三角形×3顶点）
        std::vector<uint32_t> cubeIndices = {
            0, 1, 2, 2, 1, 3, // 前面
            4, 5, 6, 6, 5, 7, // 后面
            0, 2, 4, 4, 2, 6, // 左面
            1, 3, 5, 5, 3, 7, // 右面
            0, 1, 4, 4, 1, 5, // 底面
            2, 3, 6, 6, 3, 7  // 顶面
        };

        builder.vertices = cubeVertices;
        builder.indices = cubeIndices;
        for (auto &v : builder.vertices)
        {
            v.position += offset;
        }
        return std::make_unique<NjmModel>(device, builder);
    }

    void ECSApp::loadModels()
    {
        std::cout << "--ECSApp::loadModels--" << std::endl;
        std::shared_ptr<NjmModel> cubeModel = ecsCreateCubeModel(njmDevice, {0.0f, 0.0f, 0.0f});
        cubeModel->setId(1);
        cubeModel->setName("cube");
        njmModels.emplace(cubeModel->getId(), std::move(cubeModel));

        std::string modelPath = "src/models/md1/md1.obj";
        std::string absoluteModelPath = njm::getCompletePath(modelPath);
        std::shared_ptr<NjmModel> objModel = NjmModel::createModelFromFile(njmDevice, modelPath);
        objModel->setId(2);
        objModel->setName("md1");
        njmModels.emplace(objModel->getId(), std::move(objModel));

        modelPath = "src/models/quad/quad.obj";
        absoluteModelPath = njm::getCompletePath(modelPath);
        objModel = NjmModel::createModelFromFile(njmDevice, modelPath);
        objModel->setId(3);
        objModel->setName("quad");
        njmModels.emplace(objModel->getId(), std::move(objModel));
    }

    void ECSApp::loadImgs()
    {
        const std::string &IMG_BASE_PATH = "src/texs/";

        std::string imgName = "eg1.png";
        std::string imgPath = IMG_BASE_PATH + imgName;
        std::string absoluteImgPath = njm::getCompletePath(imgPath);
        // NjmImg img1{imgName, absoluteImgPath};
        auto &img1 = std::make_shared<NjmImg>(imgName, absoluteImgPath);
        imgIds.push_back(imgName);
        njmImgs.emplace(imgName, std::move(img1));
    }

    void ECSApp::printAssetsInfo()
    {
        std::cout << "--FirstApp::printAssetsInfo--" << std::endl;
        std::cout << "njmImgs.size(): " << njmImgs.size() << std::endl;
        std::cout << "imgIds.size(): " << imgIds.size() << std::endl;
        for (auto &imgId : imgIds)
        {
            std::cout << "  " << imgId << std::endl;
        }

        std::cout << "njmModels.size(): " << njmModels.size() << std::endl;
        std::cout << "njmModels:" << std::endl;
        for (auto &model : njmModels)
        {
            std::cout << "  " << model.first << std::endl;
        }
    }

}
